package com.zf.nio.selector;

import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.Set;

/**
 * NIO 编程步骤
 * 第一步：创建 Selector 选择器
 * 第二步：创建 ServerSocketChannel 通道，并绑定监听端口
 * 第三步：设置 Channel 通道是非阻塞模式
 * 第四步：把 Channel 注册到 Socketor 选择器上，监听连接事件
 * 第五步：调用 Selector 的 select 方法（循环调用），监测通道的就绪状况
 * 第六步：调用 selectKeys 方法获取就绪 channel 集合
 * 第七步：遍历就绪 channel 集合，判断就绪事件类型，实现具体的业务操作
 * 第八步：根据业务，决定是否需要再次注册监听事件，重复执行第三步操作
 *
 * 注意：连接断开后，无论是异常断开还是正常断开，此时代表这个链接的channel一直处于readable的状态，导致每次select操作都会获取到改SocketChannel，
 * 所以断开后必须手动将将channel取消注册，正常断开连接 channel.read返回-1，异常断开连接 channel.read抛出异常
 *
 * @author Zhou Feng
 */
public class SelectorDemo02 {
    @Test
    public void serverDemo() throws IOException {
        ServerSocketChannel ssc = ServerSocketChannel.open();
        ssc.socket().bind(new InetSocketAddress("127.0.0.1", 8000));
        ssc.configureBlocking(false);
        Selector selector = Selector.open();
        // 注册 channel，并且指定感兴趣的事件是 Accept
        ssc.register(selector, SelectionKey.OP_ACCEPT);
        ByteBuffer readBuff = ByteBuffer.allocate(1024);
        ByteBuffer writeBuff = ByteBuffer.allocate(128);
        writeBuff.put("received".getBytes());
        writeBuff.flip();
        while (true) {
            int nReady = selector.select();
            Set<SelectionKey> keys = selector.selectedKeys();
            Iterator<SelectionKey> it = keys.iterator();
            while (it.hasNext()) {
                SelectionKey key = it.next();
                // 事件发生后，要么处理，要么取消（cancel），不能什么都不做，否则下次该事件仍会触发，这是因为 nio 底层使用的是水平触发
                // 每次获取 SelectionKey后，都要手动移除，否则 该SelectionKey 会一直存在于 selector.selectedKeys()的结果中

                // 如果不移除，此时 SelectionKey 虽然还在，ServerSocketChannel 也能拿到，但调用 channel.accept(); 的时候，
                // 并没有客户端真正在发起连接(上一个循环已经处理过真正的连接请求了，只是没有将这个 Key 从 SelectedKey 中移除)。
                // 所以 channel.accept(); 会返回一个 null，我们再对 null 调用 configureBlocking 方法
                it.remove();

                if (key.isAcceptable()) {
                    // 可以接受连接，创建新的连接，并且把连接注册到 selector 上，而且，
                    // 声明这个 channel 只对读操作感兴趣。
                    SocketChannel socketChannel = ssc.accept();
                    socketChannel.configureBlocking(false);
                    socketChannel.register(selector, SelectionKey.OP_READ);
                } else if (key.isReadable()) {
                    SocketChannel socketChannel = (SocketChannel) key.channel();
                    try {
                        readBuff.clear();
                        int readCount = socketChannel.read(readBuff);
                        if (readCount==-1){
                            // 注意：连接断开后，无论是异常断开还是正常断开，SocketChannel会一直处于 OP_READ状态，导致每次select操作都会获取到改SocketChannel，
                            // 所以断开后必须手动将将channel取消注册。
                            // 连接是正常断开，read返回 -1
                            // 取消注册、从注册的SelectionKey集合中移除该key
                            key.cancel();
                            // 关闭socketChannel
                            socketChannel.close();
                        } else {
                            readBuff.flip();
                            System.out.println("received : " + new String(readBuff.array()));
                            // 更新感兴趣的操作
                            key.interestOps(SelectionKey.OP_WRITE);
                        }
                    }catch (IOException e){
                        // 异常断开，抛出异常
                        // 取消注册、从注册的SelectionKey集合中移除该key
                        key.cancel();
                        // 关闭socketChannel
                        socketChannel.close();
                    }

                } else if (key.isWritable()) {
                    writeBuff.rewind();
                    SocketChannel socketChannel = (SocketChannel) key.channel();
                    socketChannel.write(writeBuff);
                    key.interestOps(SelectionKey.OP_READ);
                }
            }
        }
    }

    @Test
    public void clientDemo() throws InterruptedException {
        try {
            SocketChannel socketChannel = SocketChannel.open();
            socketChannel.connect(new InetSocketAddress("127.0.0.1", 8000));
            // 配置通道为非阻塞
            socketChannel.configureBlocking(false);
            ByteBuffer writeBuffer = ByteBuffer.allocate(32);
            ByteBuffer readBuffer = ByteBuffer.allocate(32);
            writeBuffer.put("hello".getBytes());
            writeBuffer.flip();
            while (true) {
                writeBuffer.rewind();
                // 非阻塞 模式下可能不能一次写完
                while (writeBuffer.hasRemaining()){
                    socketChannel.write(writeBuffer);
                }
                writeBuffer.flip();

                readBuffer.clear();
                // 没有数据可读也立即返回
                socketChannel.read(readBuffer);
                System.out.println(new String(readBuffer.array()));
                Thread.sleep(2000);
            }
        } catch (IOException e) {
        }
    }

}
